import os
import random
import numpy as np
from PIL import Image
import cv2
from datetime import datetime

foreground_dir = "/mnt/d/aiready/chonche/tmp1"  # 前景图路径
background_dir = "/mnt/d/aiready/chonche/v2/img_txt/imgs" # 背景图路径
background_txt = "/mnt/d/aiready/chonche/v2/img_txt/txt" # 背景图对应原始YOLO标签存放路径
output_image_dir = "generate2/img"  # 输出合成图保存路径
output_label_dir = "generate2/txt"  # 输出标签路径路径
num_foregrounds = 5 # 每张背景图随机放置的前景图数量
scale_range = (0.8, 1.5)    # 前景图缩放范围

# === 构建类别-文件映射 ===
foreground_classes = {} 
class_name_to_id = {}

# 遍历前景图目录，每个子文件夹对应一个类别
for idx, class_name in enumerate(os.listdir(foreground_dir)):
    class_path = os.path.join(foreground_dir, class_name)
    if not os.path.isdir(class_path):
        continue
    class_name_to_id[class_name] = idx  # 记录类别名称对应的编号
    # 获取该类别下的所有 png 文件
    files = [os.path.join(class_path, f) for f in os.listdir(class_path) if f.lower().endswith('.png')]
    if files:
        foreground_classes[idx] = files  # 存储到类别字典

# 检查前景是否加载成功
assert foreground_classes, "未找到前景图类别数据"

# === 打印并保存类别与编号映射 ===
print("\n类别映射如下：")
os.makedirs("generate1", exist_ok=True)  # 确保输出目录存在
class_file_path = os.path.join("generate1", "classes.txt")

# 保存类别对应关系到文件
with open(class_file_path, "w", encoding="utf-8") as f:
    for name, idx in class_name_to_id.items():
        print(f"{idx}: {name}")
        f.write(f"{idx} {name}\n")

print(f"类别信息已保存到：{class_file_path}")

# === 处理每张背景图 ===
for bg_filename in os.listdir(background_dir):
    bg_path = os.path.join(background_dir, bg_filename)
    try:
        bg_pil = Image.open(bg_path).convert("RGB")
    except:
        print(f"无法打开背景图：{bg_path}")
        continue

    bg_w, bg_h = bg_pil.size  # 背景图宽高
    yolo_labels = []          # 存储生成的 YOLO 标签

    # 随机叠加 num_foregrounds 个前景
    for _ in range(num_foregrounds):
        # 随机选类别 & 前景文件
        class_id = random.choice(list(foreground_classes.keys()))
        fg_path = random.choice(foreground_classes[class_id])
        try:
            fg_pil = Image.open(fg_path).convert("RGBA")  # 前景支持透明通道
        except:
            print(f"无法打开前景图：{fg_path}")
            continue

        # 缩放
        scale = random.uniform(*scale_range)
        new_w, new_h = int(fg_pil.width * scale), int(fg_pil.height * scale)
        fg_pil = fg_pil.resize((new_w, new_h), Image.LANCZOS)

        # 随机位置
        x = random.randint(0, max(0, bg_w - new_w))
        y = random.randint(0, max(0, bg_h - new_h))

        # 粘贴前景到背景
        bg_pil.paste(fg_pil, (x, y), fg_pil)

        # 生成 YOLO 格式标注
        x_center = (x + new_w / 2) / bg_w
        y_center = (y + new_h / 2) / bg_h
        width = new_w / bg_w
        height = new_h / bg_h
        yolo_labels.append(f"{class_id} {x_center:.8f} {y_center:.8f} {width:.8f} {height:.8f}")

    # 保存图像和标注 生成输出文件名
    basename = os.path.splitext(bg_filename)[0]

    # 获取当前时间戳
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    # 文件名加上时间戳
    output_img_path = os.path.join(output_image_dir, f"aug_{timestamp}_{basename}.png")
    output_lbl_path = os.path.join(output_label_dir, f"aug_{timestamp}_{basename}.txt")

    # 保存
    bg_cv = cv2.cvtColor(np.array(bg_pil), cv2.COLOR_RGB2BGR)
    cv2.imwrite(output_img_path, bg_cv)
    former_label_path = os.path.join(background_txt,basename+".txt")
    with open(former_label_path, "r") as f:
        lines = f.read().splitlines()
        yolo_labels= yolo_labels + lines

    with open(output_lbl_path, "w", encoding="utf-8") as f:
        f.write("\n".join(yolo_labels))

    print(f"合成图保存：{output_img_path}")
    print(f"标签保存：{output_lbl_path}")
